# 1.	数据预处理
import matplotlib.pyplot as plt
import numpy as np
import torch

# 1)	数据必须从CSV文件data-02-stock_daily.csv中加载，且该文件应以逗号作为分隔符。
data = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
print(data)
# 2)	数据需要被反转（data[::-1]），可能是为了使得时间序列数据按照时间顺序排列（如果原始数据是逆序存储的话）。
data = data[::-1]
# 3)	数据需要经过MinMaxScaler进行归一化处理。
from sklearn.preprocessing import MinMaxScaler

data = MinMaxScaler().fit_transform(data)
# 4)	序列长度（c）被设定为7，意味着每个输入序列包含7个连续的数据点，而目标值（y_data）是序列之后的那个数据点的最后一个特征值。
c = 7
x = []
y = []
for i in range(len(data) - c):
    x.append(data[i:i + c])
    y.append(data[i + c][-1])
x = torch.Tensor(x)
y = torch.Tensor(y).reshape(-1, 1)
# 2.	数据划分
# 1)	数据集需要被划分为训练集和测试集，且划分时不能进行随机打乱（shuffle=False）。
from sklearn.model_selection import train_test_split
train_x, test_x, train_y, test_y = train_test_split(x, y, shuffle=False)
# 3.	模型结构
# 1)	模型必须是一个循环神经网络（RNN），包含两层RNN层和一个全连接层。
# 2)	第一层RNN的input_size必须与输入数据的特征数量相匹配，这里由x.shape[2]决定。
# 3)	第一层RNN的hidden_size为5，第二层RNN的hidden_size为10。
# 4)	全连接层的输入特征数为10，输出特征数为1，用于预测单个值。
class RNN(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.rnn1 = torch.nn.RNN(input_size=x.shape[2],hidden_size=10,batch_first=True)
        self.rnn2 = torch.nn.RNN(input_size=10,hidden_size=5,batch_first=True)
        self.fc = torch.nn.Linear(in_features=5,out_features=1)
    def forward(self,x):
        x,_ = self.rnn1(x)
        x,_ = self.rnn2(x)
        out = self.fc(x[:,-1,:])
        return out
# 4.	模型训练
# 1)	使用均方误差（MSE）作为损失函数。
model = RNN()
loss_fn = torch.nn.MSELoss()
# 2)	使用Adam优化器进行参数优化，学习率设定为0.01。
op = torch.optim.Adam(params=model.parameters(),lr=0.01)
# 3)	训练过程中，模型需要进行1000次迭代。
# 4)	在每次迭代中，模型需要进行前向传播、损失计算、反向传播和参数更新。
loss_list = []
for i in range(1000):
    op.zero_grad()
    h = model(train_x)
    loss = loss_fn(h,train_y)
    loss.backward()
    loss_list.append(loss)
    op.step()
    print(loss)
# 5.	预测与可视化
# 1)	训练完成后，模型需要对测试集进行预测。
# 2)	预测结果需要被重塑为一维数组。
pred = model(test_x).reshape(-1,1)
print(pred)
# 3)	使用matplotlib库将预测结果（红色）和实际值（绿色）绘制在同一张图上以便可视化对比
plt.plot(test_y.data.numpy(),c='r')
plt.plot(pred.data.numpy(),c='g')
plt.show()